import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import load_model
import pickle
import argparse
from datetime import datetime

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def load_and_predict(weather_data_path, model_path='temperature_forecast_model.h5',
                     model_info_path='model_info.pkl', days_to_predict=7):
    """加载模型并预测未来温度"""
    print(f"加载模型：{model_path}")
    print(f"加载模型信息：{model_info_path}")
    print(f"加载气象数据：{weather_data_path}")

    # 1. 加载模型和相关信息
    model = load_model(model_path)
    with open(model_info_path, 'rb') as f:
        model_info = pickle.load(f)

    scaler = model_info['scaler']
    feature_cols = model_info['feature_cols']
    target_idx = model_info['target_idx']
    time_step = model_info['time_step']
    all_cols = model_info['all_cols'] if 'all_cols' in model_info else feature_cols

    # 2. 加载并处理气象数据
    df = pd.read_csv(weather_data_path)
    df.columns = [
        'Maximum Temperature',
        'Total Rainfall',
        'Mean Pressure',
        'Mean Wind Speed',
        'Mean Dew Point Temperature',
        'Date'
    ]
    df['Date'] = pd.to_datetime(df['Date'])
    df.set_index('Date', inplace=True)
    df.sort_index(inplace=True)

    # 3. 添加滞后特征
    for i in range(1, 4):
        df[f'Temp_Lag_{i}'] = df['Maximum Temperature'].shift(i)

    # 添加季节性特征
    df['Month'] = df.index.month
    df['Day'] = df.index.day

    # 去掉有NaN的行
    df = df.dropna()
    print(f"处理后的数据记录数：{len(df)}")

    # 4. 提取最近的数据用于预测
    # 获取最近的time_step天数据
    recent_data = df[all_cols].values[-time_step:]
    print(f"用于预测的最近{time_step}天数据shape: {recent_data.shape}")

    # 标准化数据
    recent_data_scaled = scaler.transform(recent_data)

    # 5. 预测未来温度
    last_sequence = recent_data_scaled.reshape(1, time_step, len(all_cols))
    future_predictions = []
    future_dates = []

    for i in range(days_to_predict):
        # 预测下一天
        next_pred = model.predict(last_sequence, verbose=0)[0]
        future_predictions.append(next_pred)

        # 更新序列
        temp_seq = last_sequence[0].copy()
        temp_seq = np.roll(temp_seq, -1, axis=0)
        temp_seq[-1, target_idx] = next_pred
        last_sequence = temp_seq.reshape(1, time_step, len(all_cols))

        # 生成日期
        next_date = df.index[-1] + pd.Timedelta(days=i + 1)
        future_dates.append(next_date)

    # 6. 反向转换预测结果
    future_pred_placeholder = np.zeros((len(future_predictions), len(all_cols)))
    future_pred_placeholder[:, target_idx] = np.array(future_predictions).reshape(-1)
    future_pred_inv = scaler.inverse_transform(future_pred_placeholder)[:, target_idx]

    # 7. 可视化结果
    plt.figure(figsize=(14, 8))
    plt.plot(df.index[-30:], df['Maximum Temperature'].values[-30:], label='历史温度', color='blue')
    plt.plot(future_dates, future_pred_inv, 'r--', label='未来预测')
    plt.scatter(future_dates, future_pred_inv, color='red', s=50)
    plt.title(f'常州市未来{days_to_predict}天最高温度预测')
    plt.xlabel('日期')
    plt.ylabel('温度 (°C)')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()

    # 保存预测图
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    plt.savefig(f'prediction_{timestamp}.png')
    plt.show()

    # 8. 返回预测结果
    result_df = pd.DataFrame({
        'Date': future_dates,
        'Predicted_Temperature': future_pred_inv
    })
    result_df.set_index('Date', inplace=True)

    # 保存预测结果到CSV
    csv_file = f'prediction_{timestamp}.csv'
    result_df.to_csv(csv_file)
    print(f"预测结果已保存到 {csv_file}")

    return result_df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='温度预测')
    parser.add_argument('--data', type=str, default='../常州.csv', help='气象数据CSV文件路径')
    parser.add_argument('--model', type=str, default='temperature_forecast_model.h5', help='模型文件路径')
    parser.add_argument('--info', type=str, default='model_info.pkl', help='模型信息文件路径')
    parser.add_argument('--days', type=int, default=14, help='预测天数')

    args = parser.parse_args()

    # 使用模型进行预测
    predictions = load_and_predict(
        args.data,
        model_path=args.model,
        model_info_path=args.info,
        days_to_predict=args.days
    )

    print("\n未来预测结果:")
    print(predictions)
